package com.jd.dlc.limit;

import com.jd.dlc.domain.LimiterRule;
import com.jd.dlc.spring.SpringUtils;
import com.jd.jim.cli.Cluster;
import com.jd.jim.cli.protocol.ScriptOutputType;

import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @ClassName: JimdbLimiter
 * @Description: jimdb限流
 * @Author: MengXiangXin
 * @Date: 2022/5/9 10:56 AM
 */
public class JimdbLimiter implements Limiter {


    private String sha;

    private volatile AtomicBoolean shaLoading = new AtomicBoolean(false);
    /**
     * jimdb 客户端
     */
    private Cluster jimClient = (Cluster) SpringUtils.getBean("jimClient");

    // TODO: 2022/5/9 超时 
    public static final String SCRIPT = "local key = KEYS[1]\n" +
            "local limit = tonumber(ARGV[1])\n" +
            "local current = tonumber(redis.call('get', key) or 0)\n" +
            "if current + 1 > limit then\n" +
            "    return 0\n" +
            "else\n" +
            "    redis.call(\"INCRBY\", key,\"1\")\n" +
            "    redis.call(\"expire\", key,\"2\")\n" +
            "    return 1\n" +
            "end\n";


    @Override
    public boolean isAllowed(String limitKey, long timeout, long limit) {

        if (shaLoading.compareAndSet(false, true)) {
            sha = jimClient.scriptLoad(SCRIPT);
        }

        String key = PREFIX_LIMIT + limitKey;
        boolean allowed = (Boolean) jimClient.evalsha(sha, Arrays.asList(key), Arrays.asList(String.valueOf(limit)), true, ScriptOutputType.BOOLEAN);
        return allowed;
    }

    @Override
    public boolean refreshLimiter(LimiterRule limiterRule) {
        return false;
    }
}
